import json

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
from tqdm import tqdm

device = None


def get_device():
    '''
    获取本机的设备
    :return:
    '''
    global device
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return device


def load_json_file(file):
    '''
    加载json文件
    :param file:
    :return:
    '''
    with open(file, 'r') as f:
        return json.load(f)


def save_json_file(data, file):
    '''
    保存数据到文件
    :param data:
    :param file:
    :return:
    '''
    # 使用 json.dump() 将数据保存为 JSON 文件
    with open(file, 'w') as f:
        json.dump(data, f)


def save_model(model, model_name):
    '''
    保存模型
    :param model:
    :param model_name:
    :return:
    '''
    torch.save(model.state_dict(), model_name + '.pth')


def load_model(model, model_name):
    '''
    加载模型
    :param model:
    :param model_name:
    :return:
    '''
    model.load_state_dict(torch.load(model_name + '.pth'))


def init_parameter_normal(model, mean=0, std=1):
    '''
    初始化模型中线性层的参数
    :param model:
    :param mean:
    :param std:
    :return:
    '''

    def init_weights(m):
        if type(m) == nn.Linear or type(m) == nn.Conv2d:
            print("init weight normal")
            torch.nn.init.normal_(m.weight, mean=mean, std=std)

    model.apply(init_weights)


def init_parameter_xavier_uniform(model):
    '''
    初始化模型中线性层的参数
    :param model:
    :param mean:
    :param std:
    :return:
    '''

    def init_weights(m):
        if type(m) == nn.Linear or type(m) == nn.Conv2d:
            print("init weight xavier_uniform")
            torch.nn.init.xavier_uniform_(m.weight)

    model.apply(init_weights)


def plot_line(data, labels, save_chart_name, figsize=(8, 4)):
    '''
    画线
    :param data:
    :param labels:
    :param save_chart_name:保存的图表名字
    :param figsize: 图像大小，单位是英寸
    :return:
    '''
    x_label = list(range(len(data[0])))
    line_number = len(data)
    plt.figure(figsize=figsize)
    index = 1
    for data_item in data:
        plt.subplot(line_number, 1, index)
        plt.plot(x_label, data_item)
        plt.title(labels[index - 1])
        index += 1
    # 调整子图之间的间距
    plt.tight_layout()
    # 显示图形
    plt.savefig(save_chart_name, dpi=200)
    plt.show()
    plt.close()


def train(model, train_data, test_data, loss, optim, num_epochs, device, continue_train=False,
          cache_model_info_file="model_persistence.json", model_save_name="default_model", dtype=torch.float):
    '''
    :param model: 模型
    :param train_data:训练集
    :param test_data:测试集
    :param loss:损失函数
    :param optim: 优化器
    :param num_epochs:迭代轮次
    :param device:设备
    :param continue_train:继续训练
    :param cache_model_info_file:模型缓存信息
    :param model_save_name:模型保存名字
    :param dtype:模型的数据类型
    :return:
    '''
    if device is None:
        device = get_device()
    best_vali_loss = None
    if continue_train:
        model_cache_info = load_json_file(cache_model_info_file)
        best_vali_model_save_path = model_cache_info["best_vali_model_save_path"]
        if best_vali_model_save_path is not None:
            load_model(model, best_vali_model_save_path)
            model_save_name = best_vali_model_save_path
            best_vali_loss = model_cache_info["best_vali_loss"]
            print("继续之前的训练,最优验证集损失{best_vali_loss},模型存储位置{best_vali_model_save_path}")
        else:
            print("模型加载异常,无法加载模型")
    model.to(device=device, dtype=dtype)
    train_loss_history = []
    test_loss_history = []
    for epoch in tqdm(range(num_epochs), desc="train_epoch"):
        model.train()
        train_loss = []
        for batch, (x, y) in enumerate(tqdm(train_data, desc=f"train_batch", ascii=True)):
            optim.zero_grad()
            batch_x = x.to(device=device, dtype=dtype)
            batch_y = y.to(device=device, dtype=dtype)
            outputs = model(batch_x)
            batch_train_loss = loss(outputs, batch_y)
            train_loss.append(batch_train_loss.item())
            # 反向传播更新梯度
            batch_train_loss.backward()
            # 根据梯度更新参数
            optim.step()
        epoch_loss_value = np.array(train_loss).mean()
        train_loss_history.append(epoch_loss_value)

        test_loss = []
        for batch, (x, y) in enumerate(tqdm(test_data, desc=f"test_batch", ascii=True)):
            # 更改模型状态为验证
            model.eval()
            batch_x = x.to(device=device, dtype=dtype)
            batch_y = y.to(device=device, dtype=dtype)
            outputs = model(batch_x)
            batch_test_loss = loss(outputs, batch_y)
            test_loss.append(batch_test_loss.item())

        epoch_loss_value = np.array(test_loss).mean()
        test_loss_history.append(epoch_loss_value)

        if best_vali_loss is None or best_vali_loss > test_loss_history[-1]:
            best_vali_loss = test_loss_history[-1]
            save_model(model, model_save_name)
            save_json_file({"best_vali_model_save_path": model_save_name, "best_vali_loss": best_vali_loss},
                           cache_model_info_file)

        print(
            f"\nEpoch：{epoch} Train Loss: {train_loss_history[-1]}, Test Loss: {test_loss_history[-1]}\n")

        plot_line([train_loss_history, test_loss_history], ("train_loss", "test_loss"),
                  model_save_name + "_epoch_" + str(epoch))


def valid(model, valid_data, loss, device, dtype=torch.float):
    '''
    模型验证
    :param model: 模型
    :param valid_data: 验证集
    :param device:设备
    :param dtype:模型的数据类型
    :return: 验证集损失，预测数据，真实数据
    '''
    if device is None:
        device = get_device()

    model.to(device=device, dtype=dtype)

    valid_loss = []
    predict_result = []
    true_result = []
    for batch, (x, y) in enumerate(tqdm(valid_data, desc=f"test_batch", ascii=True)):
        # 更改模型状态为验证
        model.eval()
        batch_x = x.to(device=device, dtype=dtype)
        batch_y = y.to(device=device, dtype=dtype)
        outputs = model(batch_x)
        batch_test_loss = loss(outputs, batch_y)
        valid_loss.append(batch_test_loss.item())
        predict_result.append(outputs)
        true_result.append(batch_y)

    valid_loss_value = np.array(valid_loss).mean()
    return valid_loss_value, predict_result, true_result


def train_test():
    from torch.utils.data import Dataset, DataLoader
    import torch.optim as optim
    # 示例模型
    class SimpleModel(nn.Module):
        def __init__(self):
            super(SimpleModel, self).__init__()
            self.fc1 = nn.Linear(1, 200)
            self.fc2 = nn.Linear(200, 50)
            self.fc3 = nn.Linear(50, 1)

        def forward(self, x):
            x = torch.relu(self.fc1(x))
            x = torch.relu(self.fc2(x))
            x = self.fc3(x)
            return x

    class CustomDataset(Dataset):
        def __init__(self, data, labels):
            self.data = torch.tensor(data)
            self.labels = torch.tensor(labels)

        def __len__(self):
            return len(self.data)

        def __getitem__(self, idx):
            sample = self.data[idx]
            label = self.labels[idx]
            return sample.view(1), label.view(1)

    x = np.linspace(0, 2 * np.pi, 1000, dtype=np.float32)
    y = np.sin(x)
    train_length = int(len(x) * 0.8)
    train_loader = DataLoader(CustomDataset(x[0:train_length], y[0:train_length]), batch_size=20, shuffle=True)
    test_loader = DataLoader(CustomDataset(x[train_length:-1], y[train_length:-1]), batch_size=20)

    # 实例化模型
    model = SimpleModel()
    loss = nn.MSELoss()
    optimizer = optim.SGD(model.parameters(), lr=0.005, momentum=0.9)
    train(model, train_loader, test_loader, loss, optimizer, 100, get_device())


if __name__ == "__main__":
    x = np.linspace(0, 2 * np.pi, 100)
    y = np.sin(x)
    plot_line([y, y], ("a", "b"), figsize=(16, 16))
    train_test()
